#include "codec/ffmpeg_decoder.h"
extern "C" {
#include <libavformat/avformat.h>
#include <libavcodec/avcodec.h>
#include <libswscale/swscale.h>
}
#include <stdexcept>

struct FFmpegDecoder::Impl {
    AVFormatContext* fmt_ctx = nullptr;
    AVCodecContext* codec_ctx = nullptr;
    AVFrame* frame = nullptr;
    AVPacket* pkt = nullptr;
    SwsContext* sws_ctx = nullptr;
    int video_stream_idx = -1;
    int width = 0, height = 0;
    bool opened = false;
};

FFmpegDecoder::FFmpegDecoder() : impl_(new Impl) {
    av_register_all();
    avformat_network_init();
}

FFmpegDecoder::~FFmpegDecoder() {
    close();
    delete impl_;
}

bool FFmpegDecoder::open(const std::string& input_url) {
    close();
    auto& d = *impl_;
    if (avformat_open_input(&d.fmt_ctx, input_url.c_str(), nullptr, nullptr) < 0) return false;
    if (avformat_find_stream_info(d.fmt_ctx, nullptr) < 0) return false;
    for (unsigned i = 0; i < d.fmt_ctx->nb_streams; ++i) {
        if (d.fmt_ctx->streams[i]->codecpar->codec_type == AVMEDIA_TYPE_VIDEO) {
            d.video_stream_idx = i;
            break;
        }
    }
    if (d.video_stream_idx < 0) return false;
    AVCodecParameters* codecpar = d.fmt_ctx->streams[d.video_stream_idx]->codecpar;
    AVCodec* decoder = avcodec_find_decoder(codecpar->codec_id);
    d.codec_ctx = avcodec_alloc_context3(decoder);
    avcodec_parameters_to_context(d.codec_ctx, codecpar);
    if (avcodec_open2(d.codec_ctx, decoder, nullptr) < 0) return false;
    d.frame = av_frame_alloc();
    d.pkt = av_packet_alloc();
    d.width = d.codec_ctx->width;
    d.height = d.codec_ctx->height;
    d.sws_ctx = sws_getContext(d.width, d.height, d.codec_ctx->pix_fmt, d.width, d.height, AV_PIX_FMT_BGR24, SWS_BILINEAR, nullptr, nullptr, nullptr);
    d.opened = true;
    return true;
}

bool FFmpegDecoder::readFrame(cv::Mat& frame) {
    auto& d = *impl_;
    if (!d.opened) return false;
    while (av_read_frame(d.fmt_ctx, d.pkt) >= 0) {
        if (d.pkt->stream_index == d.video_stream_idx) {
            int ret = avcodec_send_packet(d.codec_ctx, d.pkt);
            av_packet_unref(d.pkt);
            if (ret < 0) continue;
            ret = avcodec_receive_frame(d.codec_ctx, d.frame);
            if (ret == 0) {
                cv::Mat img(d.height, d.width, CV_8UC3);
                uint8_t* dst[] = { img.data };
                int dst_linesize[] = { static_cast<int>(img.step) };
                sws_scale(d.sws_ctx, d.frame->data, d.frame->linesize, 0, d.height, dst, dst_linesize);
                frame = img;
                return true;
            }
        } else {
            av_packet_unref(d.pkt);
        }
    }
    return false;
}

void FFmpegDecoder::close() {
    auto& d = *impl_;
    if (d.frame) av_frame_free(&d.frame);
    if (d.pkt) av_packet_free(&d.pkt);
    if (d.codec_ctx) avcodec_free_context(&d.codec_ctx);
    if (d.fmt_ctx) avformat_close_input(&d.fmt_ctx);
    if (d.sws_ctx) sws_freeContext(d.sws_ctx);
    d.opened = false;
}

bool FFmpegDecoder::isOpened() const { return impl_->opened; }
